package com.econ.powercloud.ems.websocket;

import com.econ.powercloud.ems.rpc.IFeignPowerNodeGatewayClient;
import com.econ.powercloud.powernode.gateway.base.utils.CurrentUserUtil;
import com.econ.powercloud.powernode.gateway.common.constants.SysConstants;
import com.econ.powercloud.powernode.gateway.common.domain.UserDTO;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.messaging.support.NativeMessageHeaderAccessor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.util.CollectionUtils;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

import java.security.Principal;
import java.text.MessageFormat;
import java.util.List;

/**
 * websocket配置类
 *
 * @author ben on 2019/4/3.
 */
@Configuration
@EnableWebSocketMessageBroker
public class WebsocketConfig implements WebSocketMessageBrokerConfigurer {
    private static final String UID = "userId";
    private static final int NUMBER_TEN_THOUSAND = 10000;

    private final String endpointURI = "/api/v2/ws/websocket";
    private final String simpleBrokerPrefixes = "/topic";
    private final String appPrefixes = "/server";
    private final String userPrefixes = "/user";

    @Autowired
    private WebsocketObserverManager websocketObserverManager;
    @Autowired
    private IFeignPowerNodeGatewayClient iFeignPowerNodeGatewayClient;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint(endpointURI)
                .setAllowedOrigins("*")
                .withSockJS()
                .setHeartbeatTime(30000);
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        ThreadPoolTaskScheduler te = new ThreadPoolTaskScheduler();
        te.setPoolSize(1);
        te.setThreadNamePrefix("wss-heartbeat-thread-");
        te.initialize();
        //服务端广播消息的路径前缀，客户端需要相应订阅/topic/yyy这个地址的消息
        registry.enableSimpleBroker(simpleBrokerPrefixes)
        .setHeartbeatValue(new long[]{NUMBER_TEN_THOUSAND, NUMBER_TEN_THOUSAND})
        .setTaskScheduler(te);
        //客户端需要把消息发送到/message/xxx地址
        registry.setApplicationDestinationPrefixes(appPrefixes);
        registry.setUserDestinationPrefix(userPrefixes);
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {

            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    String token = this.getToken(message);
                    UserDTO userDTO = iFeignPowerNodeGatewayClient.getUserBySaToken(token);
                    if (userDTO == null) {
                        throw new RuntimeException("websocket认证失败,用户未登录.");
                    }
                    String uid = userDTO.getId();
                    accessor.setUser(new Principal() {
                        @Override
                        public String getName() {
                            return uid;
                        }
                    });
                }
                return message;
            }

            @Override
            public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
                StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
                String sessionId = accessor.getSessionId();
                StompCommand command = accessor.getCommand();
                if (StompCommand.CONNECT.equals(command)) {
                    Principal principal = accessor.getUser();
                    if (principal != null) {
                        System.out.println(MessageFormat.format("用户{0}的WebSocket连接已经连接", sessionId));
                    }
                }
                if (StompCommand.SUBSCRIBE.equals(command)) {
                    System.out.println("--------------SUBSCRIBE------------");
                    LinkedMultiValueMap nativeHeaders = (LinkedMultiValueMap) accessor.getHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS);
                    System.out.println(nativeHeaders);
                }
                //用户已经断开连接
                if (StompCommand.DISCONNECT.equals(command)) {
                    System.out.println(command.getMessageType().toString());
                    Principal principal = accessor.getUser();
                    if (principal != null) {
                        //断开连接时清除本次连接订阅点号
                        websocketObserverManager.clearDataBySessionIdAndScenes(sessionId,null);
                        System.out.println(MessageFormat.format("用户{0}的WebSocket连接已经断开", principal.getName()));
                    }
                }
            }

            private String getToken(Message<?> message) {
                try {
                    Object raw = message.getHeaders().get(SimpMessageHeaderAccessor.NATIVE_HEADERS);
                    if (!(raw instanceof LinkedMultiValueMap)) {
                        return null;
                    }
                    List<String> tokenStrs = ((LinkedMultiValueMap<String, String>) raw).get("Authorization");
                    if (CollectionUtils.isEmpty(tokenStrs)) {
                        return null;
                    }
                    return tokenStrs.get(0).replaceFirst("Bearer ", "");
                } catch (Exception e) {
                    throw new RuntimeException("websocket认证失败", e);
                }
            }
        });
    }

}
